# -*- coding: utf-8 -*-
# +
import pandas as pd
import numpy as np
import os
from math import ceil
import datetime

from config import CONFIG


# +
def find_val(val, index):
    for i in range(len(CONFIG().AQI_LIMIT[index])):
        if val == 0:
            return CONFIG().IAQI_LIMIT[0], CONFIG().IAQI_LIMIT[1], CONFIG().AQI_LIMIT[index][0], CONFIG().AQI_LIMIT[index][1]
        if val < CONFIG().AQI_LIMIT[index][i]:
            return CONFIG().IAQI_LIMIT[i-1], CONFIG().IAQI_LIMIT[i], CONFIG().AQI_LIMIT[index][i-1], CONFIG().AQI_LIMIT[index][i]
    print('[no value error]', index, val)
    return CONFIG().IAQI_LIMIT[-2], CONFIG().IAQI_LIMIT[-1], CONFIG().AQI_LIMIT[index][-2], CONFIG().AQI_LIMIT[index][-1]

def aqi_compute(df, station):
    df_temp = df.copy()
    df_temp['首要污染物'] = 0
    df_temp['MAX_AQI'] = 0
    index_names = list(CONFIG().AQI_NAME_C.values() if station=='C' else CONFIG().AQI_LIMIT.keys())
    for row in list(df_temp.index):
        IAQI = {}
        MAX_VAL, MAX_INDEX = 0, None
        for index in CONFIG().AQI_LIMIT:
            index_name = CONFIG().AQI_NAME_C[index] if station=='C' else index
            IAQI_Hi, IAQI_Lo, BP_Hi, BP_Lo = find_val(df_temp.loc[row, index_name], index)
            IAQI[index] = ceil((IAQI_Hi - IAQI_Lo) / (BP_Hi - BP_Lo) * (df_temp.loc[row, index_name] - BP_Lo) + IAQI_Lo)
            df_temp.loc[row, index_name] = IAQI[index]
            if index in ['O3最大八小时滑动平均监测浓度(μg/m³)', 'O3实测八小时滑动平均日最大值(μg/m³)'] and df_temp.loc[row, index_name] > 800:
                continue
            if IAQI[index] > MAX_VAL:
                MAX_VAL, MAX_INDEX = IAQI[index], index
        df_temp.loc[row, '首要污染物'] = MAX_INDEX
        df_temp.loc[row, 'MAX_AQI'] = MAX_VAL
    df_temp = df_temp.drop(labels=[i for i in list(df_temp) if 'Unnamed' in i], axis=1)
    df_temp = df_temp.rename(columns={i:'AQI_%s'%i for i in index_names})
    return df_temp
            
def aqi_compute_from_hours(df, station):
    time_name = '监测时间' if '监测时间' in list(df) else '实测时间'
    if time_name not in list(df): print(list(df))
    df[time_name] = pd.to_datetime(df[time_name])
    df['监测日期'] = df[time_name].apply(lambda x: datetime.datetime.strftime(x, '%Y-%m-%d'))
    df_temp = pd.DataFrame()
    features = CONFIG().LABELS#[CONFIG().RENAME_DICT[i] for i in CONFIG().SC_FEATURES[station] if CONFIG().RENAME_DICT[i] in list(children)]
    for day in list(set(df['监测日期'])):
        temp = df[df['监测日期']==day].drop_duplicates(subset=['监测日期'])
        temp = temp[['监测日期']+features]
        for fea in features:
            if 'O3_' != fea:
                temp[fea] = df[df['监测日期']==day][fea].mean()
        window_8_hours = df[df['监测日期']==day]['O3_'].rolling(8).mean().dropna()
        a = max(window_8_hours) if len(window_8_hours)!=0 else 0
        temp['O3_'] = a
        df_temp = pd.concat([df_temp, temp], axis=0)
#     print(df_temp)
    df_temp['首要污染物'] = 0
    df_temp['MAX_AQI'] = 0
    index_names = features
    for row in list(df_temp.index):
        IAQI = {}
        MAX_VAL, MAX_INDEX = 0, None
        for index in index_names:
            IAQI_Hi, IAQI_Lo, BP_Hi, BP_Lo = find_val(df_temp.loc[row, index], index)
            IAQI[index] = ceil((IAQI_Hi - IAQI_Lo) / (BP_Hi - BP_Lo) * (df_temp.loc[row, index] - BP_Lo) + IAQI_Lo)
            df_temp.loc[row, index] = IAQI[index]
            if index == 'O3' and df_temp.loc[row, index] > 800:
                continue
            if IAQI[index] > MAX_VAL:
                MAX_VAL, MAX_INDEX = IAQI[index], index
        df_temp.loc[row, '首要污染物'] = MAX_INDEX
        df_temp.loc[row, 'MAX_AQI'] = MAX_VAL
    df_temp[time_name] = df[time_name]
    df_temp = df_temp.drop(labels=[i for i in list(df_temp) if 'Unnamed' in i], axis=1)
    df_temp = df_temp.rename(columns={i:'IAQI_%s'%i for i in index_names})
#     print(df_temp)
    return df_temp



# -

def run_all_aqi():
    targets = [i for i in os.listdir(CONFIG().GOOD_DATASET_PATH) if 'SCD' in i]
    for tar in targets:
        print('running target:', tar)
        file_path = os.path.join(CONFIG().GOOD_DATASET_PATH, tar)
        df = pd.read_csv(file_path)
        df2 = aqi_compute(df, tar.split('_')[0])
        df2.to_csv(os.path.join(CONFIG().IAQI_DATASET_PATH, tar))
    print('all ok.')


if __name__ == '__main__':
#     run_all_aqi()
    pass
